import os
import time
import json
import logging
from common.core.cron import CronBase
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.cache import RedisUtil
from common.enums.apply import MusicEnum
from common.enums.wallet import WalletEnum
from common.models.users import UserWalletModel
from common.models.music import MusicRecordsModel
from common.cache.queue_cache import QueueCache
from common.chain.music_server import MusicOpenAiHkServer
from plugins.storage.driver import StorageDriver

logger = logging.getLogger(__name__)


class Command(CronBase):
    @classmethod
    async def run(cls, **kwargs):
        try:
            # 读取任务
            record_id = await QueueCache.queue_pop(QueueCache.MUSIC_JOB)
            if not record_id:
                return

            # 查询任务
            record = await MusicRecordsModel.filter(id=record_id).first()
            if not record or record.status != MusicEnum.STATUS_ING:
                return

            # 处理响应
            if record.engine == "suno":
                await cls.handle_suno_check(record)
            else:
                record.status = MusicEnum.STATUS_FAIL
                record.fail_reason = f"No processing channel {record.engine}"
                await record.save()
        except Exception as e:
            logger.error(f"handle_music error: {str(e)}")

    @classmethod
    async def handle_suno_check(cls, record):
        """
        Suno音乐回调处理

        Document:
            https://openai-hk.com/docs/lab/suno-v3.html

        Author:
            zero
        """
        cover_dir_path = ""
        audio_dir_path = ""
        response = {}
        try:
            # 限流拦截
            lock_key: str = f"prevent:music_{str(record.id)}"
            lock_ack: bool = await RedisUtil.set(lock_key, str(record.id), 20, True)
            if not lock_ack:
                await QueueCache.queue_push(QueueCache.MUSIC_JOB, record.id)
                return False

            # 获取生成结果
            server = MusicOpenAiHkServer()
            response = await server.fetch(record.task_id)

            # 获取结果异常
            if not isinstance(response, list):
                error = response.get("error") or {}
                await MusicRecordsModel.filter(id=record.id).update(
                    status=int(MusicEnum.STATUS_FAIL),
                    fail_reason=error.get("message", "未知错误"),
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                    change_amount=record.use_points,
                    project="AI音乐",
                    remarks=error.get("message", "未知错误")
                )
                return None

            # 处理生成结果 [streaming=进行中, complete=完成]
            if response[0].get("status") == "complete":
                # 成功
                cover_url: str = response[0]["image_large_url"]
                audio_url: str = response[0]["audio_url"]
                duration: int = int(response[0]["metadata"].get("duration", 0))
                lyric: str = response[0]["metadata"].get("gpt_description_prompt", "")
                tags: str = response[0]["metadata"].get("tags", "")

                # 存储路径
                filename: str = ToolsUtil.make_md5_str(str(audio_url) + str(time.time()))
                curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
                cover_dir: str = f"storage/applies/music/{curr_now}/{filename}.jpeg"
                audio_dir: str = f"storage/applies/music/{curr_now}/{filename}.mp3"
                cover_dir_path: str = UrlUtil.to_root_path(path=cover_dir)
                audio_dir_path: str = UrlUtil.to_root_path(path=audio_dir)

                # 下载封面
                await ToolsUtil.download_file(url=cover_url, save_path=cover_dir_path.strip())

                # 下载音频
                await ToolsUtil.download_file(url=audio_url, save_path=audio_dir_path.strip())

                # 视频上传OSS
                engine: str = await UrlUtil.get_storage_engine()
                if engine != "local":
                    await StorageDriver.push(path=cover_dir_path, key=cover_dir.strip())
                    await StorageDriver.push(path=audio_dir_path, key=audio_dir.strip())

                # 更新记录
                await MusicRecordsModel.filter(id=record.id).update(
                    status=MusicEnum.STATUS_YES,
                    cover_url=str(cover_dir).strip(),
                    audio_url=str(audio_dir).strip(),
                    duration=duration,
                    lyric=lyric,
                    tags=tags,
                    fail_reason="",
                    notify_snap=""
                )
            elif response[0].get("status") == "streaming":
                # 进行中
                await QueueCache.queue_push(QueueCache.MUSIC_JOB, record.id)
        except Exception as e:
            if record.retry_count < 2:
                record.retry_count += 1
                record.start_time = int(time.time())
                await record.save()
                await QueueCache.queue_push(QueueCache.MUSIC_JOB, record.id)
            else:
                logger.error("Suno音乐回调失败: [%s]", str(e))
                await MusicRecordsModel.filter(id=record.id).update(
                    status=MusicEnum.STATUS_FAIL,
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    fail_reason=str(e),
                    update_time=int(time.time())
                )
                if record.use_points > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record.user_id,
                        change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                        change_amount=record.use_points,
                        project="AI音乐",
                        remarks=str(e)
                    )
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if cover_dir_path and os.path.exists(cover_dir_path):
                    os.unlink(cover_dir_path)
                if audio_dir_path and os.path.exists(audio_dir_path):
                    os.unlink(audio_dir_path)
